import sys
import os
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, mean_absolute_error, mean_squared_log_error
from DataAnalysis import TRAIN_DATA_PATH, SUBMISSION_PATH, TEST_DATA_PATH
import matplotlib.pyplot as plt
from bayes_opt import BayesianOptimization
from sklearn.utils.validation import column_or_1d
from sklearn.model_selection import cross_val_score
from tqdm import tqdm
import json
from imblearn.over_sampling import RandomOverSampler
from imblearn.under_sampling import RandomUnderSampler


class DataLoader:
    def __init__(self, is_train=True):
        self.dataset = pd.read_csv(TRAIN_DATA_PATH if is_train else TEST_DATA_PATH)
        self.keys = self.dataset.keys()
        if is_train:
            self.dataset = self.dataset.sample(frac=1, random_state=1)

        self.train_y, self.train_x = dataframe2array(self.dataset)
        print(self.train_y.shape)
        print(self.train_x.shape)

    def load_data(self):
        return self.train_y, self.train_x


def dataframe2array(data_set: pd.DataFrame):
    x = []
    y = []
    for key in data_set.keys():
        if key == 'id':
            continue
        elif key != 'price_doc':
            x.append(data_set[key].tolist())
        else:
            y.append(data_set[key].tolist())
    y = np.array(y).T
    y = y.flatten()
    x = np.array(x).T
    return y, x


def do_train(n_estimators, min_samples_leaf, min_samples_split, max_features, max_depth):
    val = cross_val_score(
        RandomForestRegressor(n_estimators=int(n_estimators),
                              min_samples_leaf=int(min_samples_leaf),
                              min_samples_split=int(min_samples_split),
                              max_features=min(max_features, 0.99999),
                              max_depth=int(max_depth),
                              random_state=2, n_jobs=-1
                              ),
        x, y, scoring='neg_mean_squared_log_error', cv=5
    ).mean()

    return val


def findparam():
    rf_bo = BayesianOptimization(
        do_train,
        {
            'n_estimators': (50, 700),
            'min_samples_split': (15, 200),
            'min_samples_leaf': (30, 200),
            'max_features': (0.01, 0.5),
            'max_depth': (5, 90)
        }
    )
    rf_bo.maximize(n_iter=30)
    _best = rf_bo.max
    print(json.dumps(_best))
    return _best


def liner_plot(ans, predictions, figure_name=''):
    plt.figure(figure_name)
    plt.scatter(ans, predictions)
    plt.xlabel('True Values ')
    plt.ylabel('Predictions ')
    plt.axis('equal')
    plt.axis('square')
    plt.xlim([0, plt.xlim()[1]])
    plt.ylim([0, plt.ylim()[1]])
    _ = plt.plot([-1e7, 1e8], [-1e7, 1e8])


def show_pred():
    """
    用来看回归效果
    """
    loader = DataLoader()
    all_y, all_x = loader.load_data()
    train_l = len(all_x) * 0.8
    train_x = all_x[:round(train_l)]
    train_y = all_y[:round(train_l)]
    test_x = all_x[round(train_l):]
    test_y = all_y[round(train_l):]
    model.fit(X=train_x, y=train_y)
    pred_y = model.predict(X=test_x)
    mse = mean_squared_error(test_y, pred_y)
    mae = mean_absolute_error(test_y, pred_y)
    rmse = np.sqrt(mse)
    msle = mean_squared_log_error(test_y, pred_y)
    rmsle = np.sqrt(msle)
    print('MSE%.4E\tMAE%.4E\tRMSE%.4E\tMSLE%.4E\tRMSLE%.4E' % (mse, mae, rmse, msle, rmsle))
    liner_plot(ans=test_y, predictions=pred_y, figure_name='pred')


def pred_submission():
    model.fit(X=x, y=y)

    test_dataloader = DataLoader(False)
    test_y, test_x = test_dataloader.load_data()
    pred = model.predict(test_x)
    with open(SUBMISSION_PATH, 'w') as subf:
        for num in tqdm(pred):
            # print('{}'.format(round(num)))
            subf.write('{}\n'.format(int(num)))
    print('finish')


if __name__ == '__main__':
    """"{"max_depth": 63.93905332700388, 
    "max_features": 0.4066899800540083, 
    "min_samples_leaf": 47.02970467525204, 
    "min_samples_split": 84.0272158616413, 
    "n_estimators": 559.7284081689106}"""

    dataloader = DataLoader()
    y, x = dataloader.load_data()
    best = findparam()
    params = best['params']
    loss = np.sqrt(-best['target'])
    print('origin loss %.4E' % loss)
    model = RandomForestRegressor(n_estimators=int(params['n_estimators']),
                                  min_samples_leaf=int(params['min_samples_leaf']),
                                  min_samples_split=int(params['min_samples_split']),
                                  max_features=params['max_features'],
                                  max_depth=int(params['max_depth']),
                                  random_state=2, n_jobs=-1
                                  )

    show_pred()
    pred_submission()
    plt.show()
